package top.haijunit.work.configurer;

import lombok.NonNull;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.HandshakeInterceptor;
import top.haijunit.work.service.WebSocketHandler;

import java.util.Map;

/**
 * 【说明类功能】
 *
 * @author zhanghaijun
 * @since 2024-09-20
 */
@Configuration
@EnableWebSocket
@RequiredArgsConstructor
public class WebSocketConfig implements WebSocketConfigurer {

    private final WebSocketHandler myWebSocketHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(myWebSocketHandler, "/ws").setAllowedOrigins("*").addInterceptors(new HandshakeInterceptor() {
            @Override
            public boolean beforeHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response, @NonNull org.springframework.web.socket.WebSocketHandler wsHandler, @NonNull Map<String, Object> attributes) throws Exception {
                // 从请求中获取查询参数并存入attributes中
                String queryString = request.getURI().getQuery();
                if (queryString != null) {
                    String[] queryParams = queryString.split("&");
                    for (String queryParam : queryParams) {
                        String[] keyValue = queryParam.split("=");
                        if (keyValue.length == 2) {
                            attributes.put(keyValue[0], keyValue[1]);
                        }
                    }
                }
                return true;
            }

            @Override
            public void afterHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response, @NonNull org.springframework.web.socket.WebSocketHandler wsHandler, Exception exception) {
                // 握手完成后执行的操作
            }
        });
    }

}
